#coding:utf8
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import time
import theano
import theano.tensor as T
from utils import generateDat
np.random.seed(123)



M,TT,dB,L = 30000, 20000, 10, 12
EqD = int(round((L+10)/2))
SNR = range(-10, 20)

title_size = 18
label_size = 16



###### RLS ######
def RLS(X,Tx):
    # 输入 测试用的X和Tx, 输出 权值和评分
    c = np.zeros( (1,L+1) );
    R_inverse = 100*np.eye(L+1)

    for k in range(TT-10):
        e = Tx[k+10+L-EqD] - c.dot( X[:,k+10]);
        filtered_infrmn_vect = R_inverse.dot(X[:,k+10]);  # (13,1)
        norm_error_power = np.conj(X[:,k+10].T).dot(filtered_infrmn_vect);
        gain_constant = 1 / (1 + norm_error_power);
        norm_filtered_infrmn_vect = gain_constant * np.conj(filtered_infrmn_vect.T);
        c = c + e * norm_filtered_infrmn_vect;
        R_inverse = R_inverse - np.conj(norm_filtered_infrmn_vect.reshape((13,1))).dot(norm_filtered_infrmn_vect.reshape((1,13)));

    sb = np.dot(c, X)
    pdvalue = sb.ravel()
    accuracy = score(pdvalue, Tx)
    return c,sb,accuracy
def score(pdvalue, Tx):
    count = 0
    for i in range(len(pdvalue)-20):
        if pdvalue[i+10].imag * Tx[i+10+L-EqD].imag >= 0 and pdvalue[i+10].real * Tx[i+10+L-EqD].real >=0:
            count += 1
    return count / (len(pdvalue)-20)
X, Tx, x = generateData(30000,20000,dB,12)
rls_weights, sb, accuracy = RLS(X,Tx)
print('RLS accuracy: {}'.format(accuracy))

print Tx.shape, sb.shape
# print plt.rcParams.keys()
plt.rcParams['font.sans-serif']=['simhei'] #用来正常显示中文标签

plt.rcParams['axes.unicode_minus']=False #用来正常显示负号
fig = plt.figure()
plt.subplots_adjust(hspace=0.4,wspace=0.3)
ax1 = fig.add_subplot(2,2,1)
ax1.scatter(x.real, x.imag)
ax1.set_title(u"接收端符号", fontsize=title_size)
ax1.set_xlabel(u"实部", fontsize=label_size)
ax1.set_ylabel(u"虚部", fontsize=label_size)
ax2 = fig.add_subplot(2,2,2)
ax2.scatter(sb[0].real, sb[0].imag)
ax2.set_title("RLS", fontsize=title_size)
ax2.set_xlabel(u"实部", fontsize=label_size)
ax2.set_ylabel(u"虚部", fontsize=label_size)

class Layer(object):
    def __init__(self, inputs, in_size, out_size, activation_function=None):
        self.W = theano.shared(np.random.normal(0,1,(in_size, out_size)))
        self.Wx_plus_b = T.dot(inputs, self.W)
        self.activation_function = activation_function
        if activation_function:
            self.outputs = self.activation_function(self.Wx_plus_b)
        else:
            self.outputs = self.Wx_plus_b

# determine the inputs
z = T.dmatrix('z')
y = T.dmatrix('y')

def score2(pdvalue, Tx):
    score = 0
    for i in range(len(pdvalue)):
        if pdvalue[i][0] * Tx[i].real >=0 and pdvalue[i][1] * Tx[i].imag >=0:
            score += 1
    return score/len(pdvalue)



def MLP(X,Tx):
    Y = X.T
    Y = np.hstack((np.real(Y),np.imag(Y))) #19990 * 26
    Txlist = np.vstack((np.real(Tx), np.imag(Tx))).T  #19990 * 2
    l1 = Layer(y, 2*L+2, 2, None)  #26入2出

    # loss function 误差函数
    cost = T.mean(T.square(l1.outputs - z))

    # compute the gradients
    gW1 = T.grad(cost, l1.W)
    # apply the gradient descent
    learning_rate = 0.05
    train = theano.function(
        inputs = [y,z],
        outputs = cost,
        updates = [(l1.W, l1.W - learning_rate * gW1)])

    # predict
    predict = theano.function(inputs=[y], outputs=l1.outputs)
    for i in range(201):
        err = train(Y, Txlist)
        if i % 50 == 0:
            pass
            # print(i, time.time()-start, err)

    testY,testTx,testx = generateData(30000,20000,dB,L)
    testY = testY.T
    testY = np.hstack((np.real(testY),np.imag(testY)))
    predictValue = predict(testY)

    return predictValue, score2(predictValue,testTx)

X, Tx, x = generateData(30000,20000,dB,L)
predictValue, accuracy = MLP(X,Tx)
print("MLP ACCURACY: {}".format(accuracy))


ax3 = fig.add_subplot(2,2,3)
predictValue = predictValue.T
ax3.scatter(predictValue[0], predictValue[1])
ax3.set_title("MLP", fontsize=title_size)
ax3.set_xlabel(u"实部", fontsize=label_size)
ax3.set_ylabel(u"虚部", fontsize=label_size)

ax4 = fig.add_subplot(2,2,4)
ax4.scatter([-1,-1,1,1],[-1,1,-1,1])
ax4.set_title("CNN", fontsize=title_size)
ax4.set_xlabel(u"实部",fontsize=label_size)
ax4.set_ylabel(u"虚部", fontsize=label_size)

# plt.savefig('foo.png')
# plt.show()
plt.show()
